"""
@description: Script to save the final trained model, handling PEFT adapter merging if necessary.
@author: yangkw
@date: 2025-05-02
"""
# -*- coding: utf-8 -*-
import logging
import argparse
import torch
from pathlib import Path
import sys
import os

# 添加 src 目录到 Python 路径
project_root = Path(__file__).parent.parent
src_path = project_root / "src"
sys.path.insert(0, str(project_root))

from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel # 导入 PeftModel 用于加载适配器

# 使用 utils 中的 logger
from src.utils import logger

def merge_and_save(
    base_model_path: str,
    adapter_path: str,
    output_path: str,
    torch_dtype_str: str = "auto", # 添加 dtype 参数
    trust_remote_code: bool = True,
):
    """
    加载基础模型和 PEFT 适配器，合并它们，并保存到指定路径。

    Args:
        base_model_path (str): 基础预训练模型的路径 (Hugging Face Hub 或本地)。
        adapter_path (str): 训练好的 PEFT 适配器权重路径 (通常是 Trainer 的 output_dir)。
        output_path (str): 合并后模型的保存路径。
        torch_dtype_str (str): 加载基础模型时使用的 torch dtype ('auto', 'bfloat16', 'float16', 'float32')。
                               合并通常建议在较高精度下进行。
        trust_remote_code (bool): 是否信任远程代码 (加载某些模型时需要)。
    """
    logger.info("="*30 + " 开始合并 PEFT 适配器 " + "="*30)
    logger.info(f"基础模型路径: {base_model_path}")
    logger.info(f"适配器路径 (检查点): {adapter_path}")
    logger.info(f"输出路径: {output_path}")
    logger.info(f"Torch Dtype: {torch_dtype_str}")

    output_path = Path(output_path)
    output_path.mkdir(parents=True, exist_ok=True) # 确保输出目录存在

    adapter_path = Path(adapter_path)
    if not adapter_path.exists() or not (adapter_path / "adapter_config.json").exists():
         logger.error(f"指定的适配器路径 '{adapter_path}' 无效或不包含 'adapter_config.json'。请确保这是 PEFT 训练的输出目录。")
         raise FileNotFoundError(f"适配器配置 'adapter_config.json' 未在 '{adapter_path}' 中找到。")

    # --- 1. 确定加载精度 ---
    torch_dtype = None
    if torch_dtype_str == "auto":
        torch_dtype = None # 让 transformers 自动推断
    elif torch_dtype_str == "bfloat16":
        if torch.cuda.is_bf16_supported():
            torch_dtype = torch.bfloat16
        else:
            logger.warning("环境不支持 bfloat16，将使用 float32 加载基础模型进行合并。")
            torch_dtype = torch.float32
    elif torch_dtype_str == "float16":
        torch_dtype = torch.float16
    elif torch_dtype_str == "float32":
        torch_dtype = torch.float32
    else:
        logger.warning(f"未知的 torch_dtype_str: {torch_dtype_str}，将使用 'auto'。")

    logger.info(f"将使用 dtype '{torch_dtype or 'auto'}' 加载基础模型进行合并。")


    # --- 2. 加载基础模型 ---
    # 注意：合并时通常不加载量化版本的基础模型，而是加载原始精度（或至少 FP16/BF16）
    # 因为 LoRA 权重是在 (反量化后的) 计算精度上训练的。
    # 如果基础模型本身很大，加载它可能需要较多内存。
    logger.info(f"正在加载基础模型: {base_model_path}...")
    try:
        base_model = AutoModelForCausalLM.from_pretrained(
            base_model_path,
            torch_dtype=torch_dtype,
            trust_remote_code=trust_remote_code,
            # device_map="auto" # 合并通常在 CPU 或单个 GPU 上完成，避免自动分片问题
            low_cpu_mem_usage=True # 尝试减少 CPU 内存使用
        )
        logger.info("基础模型加载成功。")
    except Exception as e:
        logger.error(f"加载基础模型 '{base_model_path}' 失败: {e}")
        raise

    # --- 3. 加载 Tokenizer ---
    # Tokenizer 通常与基础模型一致
    logger.info(f"正在加载 Tokenizer: {base_model_path}...")
    try:
        tokenizer = AutoTokenizer.from_pretrained(
            base_model_path,
            trust_remote_code=trust_remote_code
        )
        # 确保设置了 pad token (如果需要)
        if tokenizer.pad_token is None:
            if tokenizer.eos_token is not None:
                tokenizer.pad_token = tokenizer.eos_token
                logger.info(f"Tokenizer 设置 pad_token = eos_token: {tokenizer.eos_token}")
            else:
                 logger.warning("Tokenizer 缺少 pad_token 和 eos_token。")
        logger.info("Tokenizer 加载成功。")
    except Exception as e:
        logger.error(f"加载 Tokenizer '{base_model_path}' 失败: {e}")
        raise

    # --- 4. 加载 PEFT 模型 (将适配器加载到基础模型上) ---
    logger.info(f"正在从 '{adapter_path}' 加载 PEFT 适配器...")
    try:
        # is_trainable=False 因为我们只是加载用于合并/推理
        peft_model = PeftModel.from_pretrained(base_model, str(adapter_path), is_trainable=False)
        logger.info("PEFT 适配器加载成功。")
    except Exception as e:
        logger.error(f"从 '{adapter_path}' 加载 PEFT 适配器失败: {e}")
        raise

    # --- 5. 合并权重 ---
    logger.info("正在合并适配器权重到基础模型...")
    try:
        # merge_and_unload() 会将 LoRA 权重合并到基础模型层，并移除 PEFT 的包装器
        merged_model = peft_model.merge_and_unload()
        logger.info("权重合并完成。")
    except Exception as e:
        logger.error(f"合并 PEFT 权重失败: {e}")
        # 可能的原因：模型结构不匹配、内存不足等
        raise

    # --- 6. 保存合并后的模型和 Tokenizer ---
    logger.info(f"正在将合并后的模型保存到: {output_path}...")
    try:
        merged_model.save_pretrained(str(output_path))
        tokenizer.save_pretrained(str(output_path))
        logger.info("合并后的模型和 Tokenizer 保存成功。")
    except Exception as e:
        logger.error(f"保存合并后的模型到 '{output_path}' 失败: {e}")
        raise

    logger.info("="*30 + " PEFT 适配器合并完成 " + "="*30)


def save_checkpoint_as_full_model(
    checkpoint_path: str,
    output_path: str,
    torch_dtype_str: str = "auto",
    trust_remote_code: bool = True,
):
    """
    加载一个完整的模型检查点 (非 PEFT) 并以标准格式保存。
    这主要用于处理未使用 PEFT 训练时保存的检查点。

    Args:
        checkpoint_path (str): 模型检查点目录的路径。
        output_path (str): 保存完整模型的路径。
        torch_dtype_str (str): 加载模型时使用的 torch dtype。
        trust_remote_code (bool): 是否信任远程代码。
    """
    logger.info("="*30 + " 保存完整模型检查点 " + "="*30)
    logger.info(f"检查点路径: {checkpoint_path}")
    logger.info(f"输出路径: {output_path}")
    logger.info(f"Torch Dtype: {torch_dtype_str}")

    output_path = Path(output_path)
    output_path.mkdir(parents=True, exist_ok=True)

    checkpoint_path = Path(checkpoint_path)
    if not checkpoint_path.exists() or not (checkpoint_path / "config.json").exists():
         logger.error(f"指定的检查点路径 '{checkpoint_path}' 无效或不包含 'config.json'。")
         raise FileNotFoundError(f"模型配置 'config.json' 未在 '{checkpoint_path}' 中找到。")

    # 确定加载精度
    torch_dtype = None
    # ... (与 merge_and_save 中相同的 dtype 处理逻辑) ...
    if torch_dtype_str == "auto": torch_dtype = None
    elif torch_dtype_str == "bfloat16": torch_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float32
    elif torch_dtype_str == "float16": torch_dtype = torch.float16
    elif torch_dtype_str == "float32": torch_dtype = torch.float32

    logger.info(f"将使用 dtype '{torch_dtype or 'auto'}' 加载检查点模型。")

    # 加载模型和 Tokenizer
    try:
        logger.info(f"正在加载模型检查点: {checkpoint_path}...")
        model = AutoModelForCausalLM.from_pretrained(
            str(checkpoint_path),
            torch_dtype=torch_dtype,
            trust_remote_code=trust_remote_code,
            low_cpu_mem_usage=True
        )
        logger.info("模型检查点加载成功。")

        logger.info(f"正在加载 Tokenizer: {checkpoint_path}...")
        tokenizer = AutoTokenizer.from_pretrained(
            str(checkpoint_path),
            trust_remote_code=trust_remote_code
        )
        logger.info("Tokenizer 加载成功。")

    except Exception as e:
        logger.error(f"加载检查点 '{checkpoint_path}' 失败: {e}")
        raise

    # 保存模型和 Tokenizer
    logger.info(f"正在将模型保存到: {output_path}...")
    try:
        model.save_pretrained(str(output_path))
        tokenizer.save_pretrained(str(output_path))
        logger.info("模型和 Tokenizer 保存成功。")
    except Exception as e:
        logger.error(f"保存模型到 '{output_path}' 失败: {e}")
        raise

    logger.info("="*30 + " 完整模型检查点保存完成 " + "="*30)


def main():
    """
    脚本主入口函数。
    解析命令行参数，判断是合并 PEFT 还是保存完整检查点，并执行相应操作。
    """
    parser = argparse.ArgumentParser(description="合并 PEFT 适配器或保存完整模型检查点")
    parser.add_argument(
        "--checkpoint_path",
        type=str,
        required=True,
        help="训练输出的检查点目录路径 (包含适配器权重或完整模型权重)。"
    )
    parser.add_argument(
        "--output_path",
        type=str,
        required=True,
        help="保存合并后模型或完整模型的路径。"
    )
    parser.add_argument(
        "--base_model_path",
        type=str,
        default=None,
        help="基础预训练模型的路径 (Hugging Face Hub 或本地)。仅在合并 PEFT 适配器时需要。"
    )
    parser.add_argument(
        "--torch_dtype",
        type=str,
        default="auto",
        choices=["auto", "bfloat16", "float16", "float32"],
        help="加载基础模型或检查点时使用的 torch dtype。"
    )
    parser.add_argument(
        "--trust_remote_code",
        action='store_true', # 默认为 False，需要时添加 --trust_remote_code
        help="是否信任远程代码 (加载某些模型时需要)。"
    )

    args = parser.parse_args()

    checkpoint_path = Path(args.checkpoint_path)

    # 判断是 PEFT 检查点还是完整模型检查点
    is_peft_checkpoint = (checkpoint_path / "adapter_config.json").exists()

    if is_peft_checkpoint:
        logger.info("检测到 PEFT 适配器检查点，将执行合并操作。")
        if not args.base_model_path:
            logger.error("合并 PEFT 适配器需要提供 --base_model_path 参数。")
            sys.exit(1)
        try:
            merge_and_save(
                base_model_path=args.base_model_path,
                adapter_path=args.checkpoint_path,
                output_path=args.output_path,
                torch_dtype_str=args.torch_dtype,
                trust_remote_code=args.trust_remote_code,
            )
            logger.info("PEFT 模型合并成功。")
        except Exception as e:
            logger.exception(f"PEFT 模型合并过程中发生错误: {e}")
            sys.exit(1)
    else:
        logger.info("未检测到 PEFT 适配器配置，将尝试作为完整模型检查点保存。")
        try:
            save_checkpoint_as_full_model(
                checkpoint_path=args.checkpoint_path,
                output_path=args.output_path,
                torch_dtype_str=args.torch_dtype,
                trust_remote_code=args.trust_remote_code,
            )
            logger.info("完整模型检查点保存成功。")
        except Exception as e:
            logger.exception(f"保存完整模型检查点过程中发生错误: {e}")
            sys.exit(1)

if __name__ == "__main__":
    main()

# 如何使用:
#
# 1. 合并 PEFT 适配器:
#    假设你的基础模型是 `meta-llama/Llama-2-7b-hf`，
#    训练输出（包含适配器）在 `./llm-finetune/output/my_finetuned_model/checkpoint-1000`，
#    你想将合并后的模型保存在 `./merged_llama_lora`：
#
#    python scripts/save_merged_model.py \
#        --checkpoint_path ./llm-finetune/output/my_finetuned_model/checkpoint-1000 \
#        --base_model_path meta-llama/Llama-2-7b-hf \
#        --output_path ./merged_llama_lora \
#        --torch_dtype bfloat16  # 或者 float16, auto
#        # --trust_remote_code # 如果基础模型需要
#
# 2. 保存完整模型检查点 (未使用 PEFT):
#    假设你的训练检查点在 `./llm-finetune/output/my_full_finetuned_model/checkpoint-500`，
#    你想将其保存为标准格式到 `./final_full_model`：
#
#    python scripts/save_merged_model.py \
#        --checkpoint_path ./llm-finetune/output/my_full_finetuned_model/checkpoint-500 \
#        --output_path ./final_full_model \
#        --torch_dtype auto
#        # --trust_remote_code # 如果模型需要
#
# 这个脚本现在可以处理 PEFT 合并和完整模型检查点的保存。
